import sys
from python_ai.common.xcommon import *
import tensorflow as tf

tf.random.set_seed(777)

sep('a')
a = tf.random.uniform([3, 4], 0, 100, dtype=tf.float32)
print(a)

sep('top_k 2 sorted')
v, idx = tf.math.top_k(a, 2, sorted=True)
print(v)
print(idx)

sep('a 3d')
a = tf.random.uniform([5, 3, 4], 0, 100, dtype=tf.float32)
print(a)

sep('top_k 2 sorted')
v, idx = tf.math.top_k(a, 2, sorted=True)
print(v)
print(idx)
